import pandas as pd
import numpy as np
import pandas_datareader.data as web


def boll_bands(data, ndays):
    """
    计算布林带
    :param data: 股票的df格式数据
    :param ndays: 计算使用的简单移动均线周期
    :return:
    """
    ma = pd.Series(np.round(data['Close'].rolling(ndays).mean(), 2), name='MA%s' % ndays)  # 计算nday均线
    # pandas.std() 默认是除以n-1 的，即是无偏的，如果想和numpy.std() 一样有偏，需要加上参数ddof=0
    # 此处添加ddof的原因是wind和yahoo的计算均采用的有偏值进行的计算
    std = pd.Series(np.round(data['Close'].rolling(ndays).std(ddof=0), 2))  # 计算nday标准差，有偏
    b1 = ma + (2 * std)  # 此处的2就是Standard Deviations
    B1 = pd.Series(b1, name='UpperBollingerBand')
    data = data.join(ma)  # 上边不写name 这里报错
    data = data.join(B1)

    b2 = ma - (2 * std)
    B2 = pd.Series(b2, name='LowerBollingerBand')
    data = data.join(B2)

    return data


if __name__ == '__main__':
    data = web.DataReader('^IXIC', data_source='yahoo', start='2/5/1971', end='1/14/2019')
    data = pd.DataFrame(data)  # 得到的数据中index直接就是Date
    n = 50
    NIFTY_BBANDS = boll_bands(data, n)
    print(NIFTY_BBANDS.loc[:, ['LowerBollingerBand', 'UpperBollingerBand']])
    print(NIFTY_BBANDS.columns)
